from ryu.base import app_manager  
from ryu.controller import ofp_event  
from ryu.controller.handler import MAIN_DISPATCHER, set_ev_cls  
from ryu.ofproto import ofproto_v1_0  
from ryu.lib.packet import packet, ethernet, ipv4  
  
class CustomSwitch(app_manager.RyuApp):  
    OFP_VERSIONS = [ofproto_v1_0.OFP_VERSION]  
  
    def __init__(self, *args, **kwargs):  
        super(CustomSwitch, self).__init__(*args, **kwargs)  
  
    @set_ev_cls(ofp_event.EventOFPPacketIn, MAIN_DISPATCHER)  
    def _packet_in_handler(self, ev):  
        msg = ev.msg  
        datapath = msg.datapath  
        ofproto = datapath.ofproto  
        parser = datapath.ofproto_parser  
        dpid = datapath.id  
  
        self.logger.info("Packet in DPID: %s", dpid)  # 打印 DPID  
        pkt = packet.Packet(msg.data)  
        eth = pkt.get_protocol(ethernet.ethernet)  
        ip_pkt = pkt.get_protocol(ipv4.ipv4)  
        if ip_pkt is None:  
            return  
        src_ip = ip_pkt.src  
        dst_ip = ip_pkt.dst  
  
        self.logger.info("Received packet: SRC IP: %s, DST IP: %s", src_ip, dst_ip)  
  
        # 根据源 IP 和目的 IP 来决定流表规则  
        out_port = None  
        if dpid == 1:  # Switch 1  
            if src_ip == '192.168.0.101' and dst_ip == '192.168.0.104':  # h1 -> h4  
                out_port = 4  
            elif src_ip == '192.168.0.104' and dst_ip == '192.168.0.101':  # h4 -> h1  
                out_port = 1  
            elif src_ip == '192.168.0.102' and dst_ip == '192.168.0.105':  # h2 -> h5  
                out_port = 4  
            elif src_ip == '192.168.0.105' and dst_ip == '192.168.0.102':  # h5 -> h2  
                out_port = 2  
            elif src_ip == '192.168.0.103' and dst_ip == '192.168.0.106':  # h3 -> h6  
                out_port = 4  
            elif src_ip == '192.168.0.106' and dst_ip == '192.168.0.103':  # h6 -> h3  
                out_port = 3  
        elif dpid == 2:  # Switch 2  
            if src_ip == '192.168.0.101' and dst_ip == '192.168.0.104':  # h1 -> h4  
                out_port = 1  
            elif src_ip == '192.168.0.104' and dst_ip == '192.168.0.101':  # h4 -> h1  
                out_port = 4  
            elif src_ip == '192.168.0.102' and dst_ip == '192.168.0.105':  # h2 -> h5  
                out_port = 2  
            elif src_ip == '192.168.0.105' and dst_ip == '192.168.0.102':  # h5 -> h2  
                out_port = 4  
            elif src_ip == '192.168.0.103' and dst_ip == '192.168.0.106':  # h3 -> h6  
                out_port = 3  
            elif src_ip == '192.168.0.106' and dst_ip == '192.168.0.103':  # h6 -> h3  
                out_port = 4  
  
        if out_port is not None:  
            self.logger.info("Installing flow: DPID %s, SRC IP %s -> DST IP %s, OUTPORT %s", dpid, src_ip, dst_ip, out_port)  
            actions = [parser.OFPActionOutput(out_port)]  
            # 安装流表规则  
            match = parser.OFPMatch(dl_type=eth.ethertype, nw_src=src_ip, nw_dst=dst_ip)  
            mod = parser.OFPFlowMod(datapath=datapath, priority=1, match=match, actions=actions)  
            datapath.send_msg(mod)  
            # 立即转发数据包  
            out = parser.OFPPacketOut(datapath=datapath, buffer_id=msg.buffer_id, in_port=msg.in_port, actions=actions, data=msg.data)  
            datapath.send_msg(out)  
        else:  
            self.logger.info("No matching flow found for DPID %s, SRC IP %s, DST IP %s", dpid, src_ip, dst_ip)